﻿using System;
using System.Collections.Generic;
using System.Linq;

namespace TextEmbedding.Core
{

    public class TfIdfEmbedding : ITextEmbedding
    {
        private readonly List<string[]> _documents;
        private readonly Dictionary<string, double> _idf;
        private readonly List<string> _vocab;

        public TfIdfEmbedding(List<string> docs)
        {
            _documents = docs.Select(d => ZhTokenizer.Tokenize(d)).ToList();
            _vocab = _documents.SelectMany(d => d).Distinct().ToList();
            _idf = CalculateIdf(_documents);
        }

        private Dictionary<string, double> CalculateIdf(List<string[]> docs)
        {
            int N = docs.Count;
            var df = new Dictionary<string, int>();

            foreach (var doc in docs)
            {
                foreach (var word in doc.Distinct())
                {
                    if (!df.ContainsKey(word)) df[word] = 0;
                    df[word]++;
                }
            }

            return df.ToDictionary(kv => kv.Key, kv => Math.Log((double)N / (1 + kv.Value)));
        }

        public float[] Embedding(int dimension, string text)
        {
            var vec = new float[dimension];
            var tokens = ZhTokenizer.Tokenize(text);
            var tf = new Dictionary<string, double>();

            foreach (var word in tokens)
            {
                if (!tf.ContainsKey(word)) tf[word] = 0;
                tf[word]++;
            }

            for (int i = 0; i < Math.Min(dimension, _vocab.Count); i++)
            {
                string word = _vocab[i];
                double tfValue = tf.ContainsKey(word) ? tf[word] / tokens.Length : 0;
                vec[i] = (float)(tfValue * (_idf.ContainsKey(word) ? _idf[word] : 0));
            }

            return vec;
        }
    }

}
